package com.cheems.cheemsspringaialibaba.Controller;

import com.alibaba.cloud.ai.graph.CompileConfig;
import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.RunnableConfig;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.checkpoint.config.SaverConfig;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.cheems.cheemsspringaialibaba.domain.RequestAi;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CompletableFuture;

@RestController
@RequestMapping("/routing")
public class RoutingController {

    private final CompiledGraph compiledGraph;
    @Autowired
    public RoutingController(@Qualifier("HenauRoutingGraph") StateGraph routingGraph)
            throws GraphStateException {
        SaverConfig saverConfig = SaverConfig.builder().build();
        // 编译时可设中断点
        this.compiledGraph = routingGraph
                .compile(CompileConfig.builder().saverConfig(saverConfig).build());
    }

    @PostMapping
    public Map<String, Object> routing(@RequestBody String text) {
        return compiledGraph.invoke(Map.of("inputText", text)).get().data();
    }

    //sse流式传输处理
    @PostMapping(path = "/stream", produces = "text/event-stream")
    public SseEmitter routingStream(@RequestBody RequestAi requestAi) {
        RunnableConfig cfg = RunnableConfig.builder().threadId(requestAi.getConversationId())
                .build();
        //创建sse链接
        SseEmitter sseEmitter = new SseEmitter(Long.MAX_VALUE);
        // 异步处理数据流
        new Thread(() -> {
            try {
                compiledGraph.stream(Map.of("inputText", requestAi.getQuery()), cfg)
                        .forEachAsync(node -> {
                            try {
                                // 按照SSE格式发送数据
                                sseEmitter.send(SseEmitter.event().data(node.toString()));
                            } catch (IOException e) {
                                sseEmitter.completeWithError(e);
                            }
                        })
                        .whenComplete((v, e) -> {
                            if (e != null) {
                                sseEmitter.completeWithError(e);
                            } else {
                                sseEmitter.complete();
                            }
                        });
            } catch (Exception e) {
                sseEmitter.completeWithError(e);
            }
        }).start();
        
        return sseEmitter;
    }
}